import os
import torch
from transformers import TrainingArguments
from llama_factory.train import train
from loguru import logger
from typing import Dict, Optional, Any

from config import settings

class FineTuner:
    """模型微调器，使用LLaMA-Factory框架进行模型微调"""
    
    def __init__(self):
        """初始化微调器"""
        # 确保输出目录存在
        os.makedirs(settings.FINETUNE_OUTPUT_DIR, exist_ok=True)
        
    def fine_tune(
        self,
        model_name: str,
        dataset_path: str,
        output_dir: Optional[str] = None,
        num_epochs: int = 3,
        batch_size: int = 4,
        learning_rate: float = 2e-5,
        lora_rank: int = 8,
        lora_alpha: int = 16,
        lora_dropout: float = 0.1
    ) -> Dict[str, Any]:
        """微调指定的模型
        
        Args:
            model_name: 要微调的模型名称
            dataset_path: 数据集路径
            output_dir: 输出目录
            num_epochs: 训练轮数
            batch_size: 批次大小
            learning_rate: 学习率
            lora_rank: LoRA的秩
            lora_alpha: LoRA的alpha值
            lora_dropout: LoRA的dropout率
        
        Returns:
            微调结果字典
        """
        logger.info(f"开始微调模型: {model_name}")
        
        # 确定输出目录
        if output_dir is None:
            # 根据模型名称和时间戳生成输出目录
            import time
            timestamp = time.strftime("%Y%m%d_%H%M%S")
            output_dir = os.path.join(settings.FINETUNE_OUTPUT_DIR, f"{model_name}_ft_{timestamp}")
        
        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)
        
        # 检查数据集路径是否存在
        if not os.path.exists(dataset_path):
            raise FileNotFoundError(f"数据集路径不存在: {dataset_path}")
        
        try:
            # 设置微调参数
            # 注意：这里使用的是LLaMA-Factory的配置格式
            finetune_args = {
                "model_name_or_path": model_name if os.path.exists(model_name) else settings.HF_MODEL_PATH,
                "dataset_dir": dataset_path,
                "output_dir": output_dir,
                "overwrite_output_dir": True,
                
                # 训练参数
                "num_train_epochs": num_epochs,
                "per_device_train_batch_size": batch_size,
                "gradient_accumulation_steps": 1,
                "learning_rate": learning_rate,
                "lr_scheduler_type": "cosine",
                "warmup_ratio": 0.03,
                
                # 优化器参数
                "optim": "adamw_torch",
                "weight_decay": 0.01,
                "max_grad_norm": 1.0,
                
                # LoRA参数
                "use_lora": True,
                "lora_rank": lora_rank,
                "lora_alpha": lora_alpha,
                "lora_dropout": lora_dropout,
                "lora_target": "all",  # 应用到所有线性层
                
                # 设备参数
                "fp16": True if torch.cuda.is_available() else False,
                "bf16": False,
                "tf32": False,
                "device_map": "auto",
                
                # 日志参数
                "logging_steps": 10,
                "save_steps": 100,
                "save_total_limit": 3,
                
                # 评估参数
                "do_eval": False,  # 可以根据需要启用
                "evaluation_strategy": "no",
                
                # 其他参数
                "gradient_checkpointing": True,
                "torch_compile": False,
                "report_to": ["tensorboard"],
            }
            
            logger.info(f"微调配置: {finetune_args}")
            
            # 调用LLaMA-Factory的train函数进行微调
            # 注意：LLaMA-Factory的API可能会有变化，请根据实际版本调整
            train_result = train(**finetune_args)
            
            logger.info(f"模型微调完成，结果保存在: {output_dir}")
            
            return {
                "status": "success",
                "model_name": model_name,
                "output_dir": output_dir,
                "train_loss": train_result.get("train_loss", None),
                "epochs": num_epochs,
                "batch_size": batch_size,
                "learning_rate": learning_rate
            }
            
        except Exception as e:
            logger.error(f"模型微调失败: {str(e)}")
            raise
    
    def merge_lora_weights(
        self,
        base_model_path: str,
        lora_path: str,
        output_path: str
    ) -> Dict[str, Any]:
        """合并LoRA权重到基础模型
        
        Args:
            base_model_path: 基础模型路径
            lora_path: LoRA权重路径
            output_path: 合并后模型的输出路径
        
        Returns:
            合并结果字典
        """
        logger.info(f"开始合并LoRA权重: {lora_path} -> {base_model_path}")
        
        try:
            from transformers import AutoModelForCausalLM, AutoTokenizer
            from peft import PeftModel
            import torch
            
            # 加载基础模型
            base_model = AutoModelForCausalLM.from_pretrained(
                base_model_path,
                cache_dir=settings.HF_CACHE_DIR,
                torch_dtype=torch.float16,
                device_map="auto",
                trust_remote_code=True
            )
            
            # 加载LoRA权重
            model = PeftModel.from_pretrained(base_model, lora_path)
            
            # 合并权重
            model = model.merge_and_unload()
            
            # 加载tokenizer
            tokenizer = AutoTokenizer.from_pretrained(
                base_model_path,
                cache_dir=settings.HF_CACHE_DIR,
                trust_remote_code=True
            )
            
            # 保存合并后的模型
            os.makedirs(output_path, exist_ok=True)
            model.save_pretrained(output_path, safe_serialization=True)
            tokenizer.save_pretrained(output_path)
            
            logger.info(f"LoRA权重合并完成，结果保存在: {output_path}")
            
            return {
                "status": "success",
                "base_model": base_model_path,
                "lora_path": lora_path,
                "output_path": output_path
            }
            
        except Exception as e:
            logger.error(f"LoRA权重合并失败: {str(e)}")
            raise
    
    def create_dataset_template(
        self,
        output_path: str,
        num_samples: int = 5
    ) -> Dict[str, Any]:
        """创建微调数据集模板
        
        Args:
            output_path: 输出路径
            num_samples: 示例数量
        
        Returns:
            创建结果字典
        """
        import json
        
        # 创建数据集模板
        dataset_template = [
            {
                "instruction": "请解释什么是人工智能？",
                "input": "",
                "output": "人工智能（Artificial Intelligence，简称AI）是计算机科学的一个分支，旨在开发能够执行通常需要人类智能的任务的机器。这些任务包括学习、推理、解决问题、理解自然语言、识别图像和声音、做出决策等。"
            },
            {
                "instruction": "请计算123乘以456的结果。",
                "input": "",
                "output": "123 × 456 = 56088"
            },
            {
                "instruction": "请将以下英文翻译成中文。",
                "input": "Hello, how are you today?",
                "output": "你好，你今天怎么样？"
            },
            {
                "instruction": "请写一首关于春天的短诗。",
                "input": "",
                "output": "春光明媚百花开，
燕舞莺歌入画来。
万物复苏生机现，
人间最美是春回。"
            },
            {
                "instruction": "请解释什么是机器学习。",
                "input": "",
                "output": "机器学习是人工智能的一个分支，它赋予计算机系统从数据中学习和改进的能力，而无需明确编程。机器学习算法使用样本数据（称为训练数据）来构建模型，用于做出预测或决策。"
            }
        ]
        
        # 截取指定数量的示例
        dataset_template = dataset_template[:num_samples]
        
        # 保存到文件
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(dataset_template, f, ensure_ascii=False, indent=2)
        
        logger.info(f"数据集模板已创建: {output_path}")
        
        return {
            "status": "success",
            "output_path": output_path,
            "num_samples": num_samples
        }